import os, base64
# 创建临时目录
if not os.path.exists("tmp"):
    os.makedirs("tmp")

from uuid import uuid4
from pydantic import BaseModel

import fastapi

from utils.asr import WhisperASR
from utils.llm import OllamaChat
from utils.tts import EdgeTTS
from utils.utils import gen_random_job_id, audio_to_base64

app = fastapi.FastAPI()

asr_model = WhisperASR()
llm_model = OllamaChat()
tts_model = EdgeTTS()

class AudioItem(BaseModel):
    audio_path: str

class LLMItem(BaseModel):
    text: str
    message: list[dict[str, str]] | None = []
    system_message: str | None = ""

class TextItem(BaseModel):
    text: str

@app.post("/asr")
async def ocr(audio_item: AudioItem):
    audio_path = audio_item.audio_path
    if ("http" not in audio_path) and (not os.path.exists(audio_path)) and (len(audio_path) < 100):
        return {"code": -1, "audio": audio_path, "meg": "Audio Not Exists."}
    content = asr_model.audio2txt(audio_path)
    return {"code": 0, "content": content}

@app.post("/llm")
async def llm(llm_item: LLMItem):
    text = llm_item.text
    if (text is None) or (text == ""):
        return {"code": -1, "meg": "Text is empty."}
    return_text = llm_model.chat(text, message=llm_item.message, system_message=llm_item.system_message)
    return {"code": 0, "content": return_text}

@app.post("/tts")
async def tts(text_item: TextItem):
    text = text_item.text
    if (text is None) or (text == ""):
        return {"code": -1, "meg": "Text is empty."}
    audio_path = f"tmp/{gen_random_job_id()}.mp3"
    await tts_model(text, audio_path)
    audio_base64 = audio_to_base64(audio_path)
    os.remove(audio_path)
    return {"code": 0, "content": audio_base64}
